/*
    Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file 

    Definitions and explicit instatiations of the aux_grid class class
    used for making images of auxiliary quantities. */

#include <CCfits/CCfits>
#include "aux_grid.h"
#include "grid.h"
#include "arepo_grid.h"
#include "arepo_emission.h"
#include "blitz-fits.h"

#ifdef WITH_AREPO
extern "C" {
#include "allvars.h"
}
#endif

using namespace std;

/** Constructor loads cell data from a FITS file.  (The grid structure
    is part of the base grid object and has already been loaded at
    this point.) */
template <typename grid_type,
	  template<typename> class sampling_policy,
	  typename rng_policy> 
mcrx::aux_grid<grid_type, sampling_policy, rng_policy>::
aux_grid (boost::shared_ptr<T_grid_impl> g,
	  CCfits::ExtHDU& structure_hdu,
	  CCfits::ExtHDU& data_hdu,
	  bool normalize_for_emission) :
  g_(g),
  emission_collection<generic_chromatic_policy<aux_pars_type>, 
  sampling_policy, rng_policy>()
{
  load(g, structure_hdu, data_hdu, normalize_for_emission);
}

template <typename grid_type,
	  template<typename> class sampling_policy,
	  typename rng_policy> 
template <typename T>
void
mcrx::aux_grid<grid_type, sampling_policy, rng_policy>::
load(boost::shared_ptr<adaptive_grid<T> > g,
     CCfits::ExtHDU& structure_hdu,
     CCfits::ExtHDU& data_hdu,
     bool normalize_for_emission)
{
  using namespace CCfits;

  const int this_task = mpi_rank();

  adaptive_grid<typename T_grid_impl::T_data>& gref =
    dynamic_cast<adaptive_grid<typename T_grid_impl::T_data>&>(*g_);

  const pair<size_t, size_t> domain = gref.domain_index();

  cout << "Task " << mpi_rank() << " reading grid cell data for "
       << domain.second << " cells starting at " << domain.first << endl;
  
  // The grid structure already has to have been loaded upon
  // constructor of the grid_base object, so it's reasonable to expect
  // us to already have an open FITS file at this point.
  structure_hdu.readKey("lengthunit", units_ ["length"]);

  Column& c_m_g = data_hdu.column("mass_gas");
  Column& c_m_metals = data_hdu.column("mass_metals");
  Column& c_sfr = data_hdu.column("SFR");
  Column& c_energy  = data_hdu.column("gas_temp_m");

  // read mass unit (length unit comes from the structure hdu).
  units_ ["mass"] = c_m_g.unit();

  std::vector<T_float> m_g;
  std::vector<T_float> m_metals;
  std::vector<T_float> sfr;
  std::vector<T_float> energy;
  std::vector<vec3d> p_g;

  std::vector<T_float> L_bol;
  std::vector<T_float> mass_stars;
  std::vector<T_float> m_m_s;
  std::vector<T_float> age_m;
  std::vector<T_float> age_l;

  c_m_g.read(m_g, domain.first+1, domain.first+domain.second);
  c_m_metals.read(m_metals, domain.first+1, domain.first+domain.second);
  c_sfr.read(sfr, domain.first+1, domain.first+domain.second);
  c_energy.read(energy, domain.first+1, domain.first+domain.second);

  T_float momcon=1;
  try {
    Column& c_p_g = data_hdu.column("p_gas");
    read(c_p_g, p_g, domain.first+1, domain.second);

    // we need to convert momentum, which is in units of
    // mass*length/time, to units of mass*m/s (since what we want is the
    // velocity)
    momcon = units::convert(c_p_g.unit(),c_m_g.unit()+"*m/s");
  }
  catch (Table::NoSuchColumn&) {
    // no momentum. set to zero
    p_g.assign(m_g.size(), vec3d(0,0,0));
  }

  // If the emission data is in the grid, we use it.  Otherwise, we
  // put 0 in the emission part.
  try{
    Column& c_l_bol = data_hdu.column("L_bol");
    Column& c_ms = data_hdu.column("mass_stars");
    Column& c_mms = data_hdu.column("mass_stellar_metals");
    Column& c_agem = data_hdu.column("age_m");
    Column& c_agel = data_hdu.column("age_l");
    units_ ["luminosity"] = c_l_bol.unit();
    
    c_l_bol.read(L_bol, domain.first+1, domain.first+domain.second);
    c_ms.read(mass_stars, domain.first+1, domain.first+domain.second);
    c_mms.read(m_m_s, domain.first+1, domain.first+domain.second);
    c_agem.read(age_m, domain.first+1, domain.first+domain.second);
    c_agel.read(age_l, domain.first+1, domain.first+domain.second);
  }
  catch(Table::NoSuchColumn&) {
    L_bol.assign(domain.second,0);
    mass_stars.assign(domain.second,0);
    m_m_s.assign(domain.second,0);
    age_m.assign(domain.second,0);
    age_l.assign(domain.second,0);
  }

  assert(m_g.size()==domain.second);
  assert(m_metals.size()==domain.second);
  assert(sfr.size()==domain.second);
  assert(energy.size()==domain.second);
  assert(p_g.size()==domain.second);

  // This vector is used to collect the pointers to the emitters for
  // the call to distr_.setup.
  std::vector<emission<generic_chromatic_policy<T_content>,
    rng_policy>*> pointers;
  array_1 weights(domain.second);
  
  int i = 0;
  for (iterator c = begin (); 
       c != end(); ++c, ++i) {
    // Note the normalization, the emitter contents are not the
    // quantities themselves but quantity/emission weight.  The
    // emission weight of the cells is stellar and gas mass, and we need to
    // take that out.  Final normalization factor is
    // sum(emission weight)/N_rays, which needs to be applied to the
    // cameras afterwards.

    // these are decoded in aux_emergence::save_images.
    T_content d;
    d[aux_pars_fields::mass_gas] = m_g [i];
    d[aux_pars_fields::mass_metals] = m_metals [i];
    d[aux_pars_fields::SFR] = sfr [i];
    d[aux_pars_fields::energy] = energy [i];
    d[aux_pars_fields::mass_stars] = mass_stars [i];
    d[aux_pars_fields::mass_metals_stars] = m_m_s [i];
    d[aux_pars_fields::L_bol] = L_bol [i];
    d[aux_pars_fields::age_m] = (age_m [i]==age_m [i] ? age_m[i] : 0);
    d[aux_pars_fields::age_l] = (age_l [i]==age_l [i] ? age_l[i] : 0);
    d[aux_pars_fields::FIR] = 0; // done in separate stage

    vec3d vel(p_g[i]*momcon/m_g[i]);
    if(m_g[i]==0) {
      vel=0.0;
    }

    if(normalize_for_emission) {
      const T_float emission_weight = m_g [i] + mass_stars [i];
      if(emission_weight!=0)
	d*= 1/emission_weight;
      else
	// to avoid NaN, we must set emission to zero in this case
	d=0;
      weights(i) = emission_weight;
    }

    T_emitter dd(d, c);

    // if the cell has no data member, default-construct one before assigning
    if(!c->data()) {
      c->set_data(new typename T_grid_impl::T_data() );
    }
    c->data()->get_emitter() = dd;

    // save variables for setup of emission_collection 
    pointers.push_back(&c->data()->get_emitter());
  }
  if(i!=domain.second) {
    cerr << "Error: unexpected number of cells in grid: " << i << " vs " << domain.second << endl;
    assert(0);
  }

  // init distribution
  if(normalize_for_emission)
    this->distr_.setup(pointers, weights);
}


/** This constructor loads cell data from the Arepo data
    structures. The grid structure is handled by the supplied grid
    object, all that we do here is use the SPH data to populate our
    data. */
template <typename grid_type,
	  template<typename> class sampling_policy,
	  typename rng_policy> 
mcrx::aux_grid<grid_type, sampling_policy, rng_policy>::
aux_grid (boost::shared_ptr<T_grid_impl> g,
	  bool normalize_for_emission) :
  g_(g),
  emission_collection<generic_chromatic_policy<aux_pars_type>, 
  sampling_policy, rng_policy>()
{
  load(g, normalize_for_emission);
}

template <typename grid_type,
	  template<typename> class sampling_policy,
	  typename rng_policy> 
template <typename T>
void
mcrx::aux_grid<grid_type, sampling_policy, rng_policy>::
load(boost::shared_ptr<arepo_grid<T> > g,
     bool normalize_for_emission)
{
#ifdef WITH_AREPO
  const int nc = this->n_cells ();
  assert(nc==N_gas);
  cout << "Creating cell data from Arepo data for " << nc << " cells. " << endl;

  // First copy units from the arepo units
  units_["length"] = arepo::arepo_units.get("length");
  units_["time"] = arepo::arepo_units.get("time");
  units_["mass"] = arepo::arepo_units.get("mass");

  // to convert arepo density to our units. (convert call kind of
  // unnecessary since we know the units are the same from above, but
  // done for clarity)
  const T_float convert_mass = 
    arepo::mcon*units::convert(arepo::arepo_units.get("mass"),
			       units_.get("mass"));

  // to convert sfr to our units.
  const T_float convert_massrate = 
    arepo::mcon/arepo::tcon*
    units::convert(arepo::arepo_units.get("mass")+"/"+
		   arepo::arepo_units.get("time"),
		   units_.get("mass")+"/"+units_.get("time"));

  // This vector is used to collect the pointers to the emitters for
  // the call to distr_.setup.
  std::vector<emission<generic_chromatic_policy<T_content>,
    rng_policy>*> pointers;
  array_1 weights(nc);
  
  int i = 0;
  for (iterator c = begin (); c != end(); ++c, ++i) {
    // Note the normalization, the emitter contents are not the
    // quantities themselves but quantity/emission weight.  The
    // emission weight of the cells is stellar and gas mass, and we need to
    // take that out.  Final normalization factor is
    // sum(emission weight)/N_rays, which needs to be applied to the
    // cameras afterwards.

    // these are decoded in aux_emergence::save_images.
    T_content d;
    d[aux_pars_fields::mass_gas] = 
      SphP[i].Density*SphP[i].Volume*convert_mass;

    assert(P[i].Metallicity==SphP[i].Metallicity);
    d[aux_pars_fields::mass_metals] = 
      d[aux_pars_fields::mass_gas]*SphP[i].Metallicity;

    d[aux_pars_fields::SFR] = SphP[i].Sfr*convert_massrate;

    // use Arepo to calculate temperature directly (hopefully it outputs in K...)
    T_float ne=SphP[i].Ne;
    const double T= convert_u_to_temp(SphP[i].Utherm, SphP [i].Density, &ne);
    d[aux_pars_fields::energy] = T*d[aux_pars_fields::mass_gas];

    // all of these are from particles, not the grid.
    d[aux_pars_fields::mass_stars] = 0;
    d[aux_pars_fields::mass_metals_stars] = 0;
    d[aux_pars_fields::L_bol] = 0;
    d[aux_pars_fields::age_m] = 0;
    d[aux_pars_fields::age_l] = 0;
    d[aux_pars_fields::FIR] = 0; 
   
    if(normalize_for_emission) {
      const T_float emission_weight = 
	d[aux_pars_fields::mass_gas] + d[aux_pars_fields::mass_stars];
      if(emission_weight!=0)
	d*= 1/emission_weight;
      else
	// to avoid NaN, we must set emission to zero in this case
	d=0;
      weights(i) = emission_weight;
    }
    T_emitter dd(d, c);

    // if the cell has no data member, default-construct one before assigning
    if(!c->data()) {
      c->set_data(new typename T_grid_impl::T_data() );
    }
    c->data()->get_emitter() = dd;

    // save variables for setup of emission_collection 
    pointers.push_back(&c->data()->get_emitter());
  }

  // init distribution
  if(normalize_for_emission)
    this->distr_.setup(pointers, weights);

  // this call goes to the emission_collection class and makes it
  // possible to do intelligent stuff after the individual emitters
  // have been set up
  T_emitter::post_setup(*g_);

#else
  cerr << "Not compiled with Arepo support" << endl;
  exit(1);
#endif
}

namespace mcrx {
template class aux_grid<
  adaptive_grid<
    cell_data<
      typename emitter_types<adaptive_grid, 
			     generic_chromatic_policy<aux_pars_type>,
			     local_random>::T_emitter,
      absorber<aux_pars_type> > >,
  cumulative_sampling, local_random>;

#ifdef WITH_AREPO
template class aux_grid<
  arepo_grid<
    cell_data<
      typename emitter_types<arepo_grid, 
			     generic_chromatic_policy<aux_pars_type>,
			     local_random>::T_emitter,
      absorber<aux_pars_type> > >,
  cumulative_sampling, local_random>;
#endif
}

